Conversation
- FlashMLA sparse: admit head_size 512 (bf16 only) and skip the concat_mla_q call when the rope part is empty - MLA prefill: whitelist the GLM5Next NoPE dims (256, 0, 256) - cmake: pin vllm-flash-attn and FlashMLA to the NoPE support commits Signed-off-by: JaredforReal <w13431838023@gmail.com> Co-authored-by: Kimi Code <noreply@moonshot.cn>
… MQA GLM5Next is rope-free MLA (qk_nope_head_dim=256, qk_rope_head_dim=0, v_head_dim=256, kv_lora_rank=512), so head_size is 512. Two fixes: - Replace the unconditional head_size==512 rejection with a merged predicate that admits a rope-free bf16/auto cache on SM90, and rejects everything else (plain fp8, SM100 bf16, rope-carrying 512, and the quantized DS-MLA formats, which are served by the zero-padded 576/656B envelope wired up separately). Restricting the bf16 arm to SM90 keeps the SM100 priority order unchanged (FlashInfer stays the default for >16 heads). - Port the FlashAttention sparse zero-rope fix: when qk_rope_head_dim is 0, feed FA3 a zero 64-wide q_rope and slice k_cache from v_cache, instead of handing it an empty head-dim tensor. Add CPU tests pinning the selection contract and the SM100 order. Signed-off-by: Leoyzen <leoyzen@gmail.com>
JaredforReal
force-pushed
the
wire_sm90_kernels
branch
from
September 15, 2026 12:09
8e0b383 to
e7d59e7
Compare
Owner
|
@Leoyzen hey thanks for your great contributions |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Two follow-ups on this branch, kept as a separate commit so they can be reviewed independently.
The
head_size == 512case insupports_combinationnow has a real guard. Adding 512 toget_supported_head_sizes()unconditionally reorders the SM100 list: with more than 16 heads (TP<=2) the default flips fromFLASHINFER_MLA_SPARSEtoFLASHMLA_SPARSE, and that path has no NoPE-512 entry. The bf16 arm is therefore limited to SM90, which leaves the SM100 order untouched. Plain fp8, SM100 bf16 and a rope-carrying 512 model are rejected; the quantized DS-MLA formats are rejected here too and handled separately in vllm-project#55543.Second, the zero-rope fix from vllm-project#56665 is ported into
flashattn_mla_sparse.py. vllm-project#56665 was closed pointing at vllm-project#55385, but this branch never touched that file, soforward_mqastill hands FA3 an empty head-dim tensor for rope-free models.Added
tests/v1/attention/test_flashmla_nope_sm90_backend_selection.py(CPU-only): the selection matrix plus a guard that the SM100 bf16 order is unchanged.pytest tests/v1/attention/test_flashmla_nope_sm90_backend_selection.py-> 28 passed, 3 skipped (GPU-only).The bf16-512 FlashMLA decode path has not been run on hardware here; it needs an SM90 smoke before this is relied on.